# -*- coding: utf-8 -*-
"""
Created on Sun Aug 18 16:13:18 2019

@author: gby
"""
import networkx as nx
from networkx import readwrite
from networkx.readwrite import json_graph
import json
import matplotlib.pyplot as plt
import pandas as pd
import numpy as np

name = "twitch_es"

# 导入graph：
options = {'node_color': 'black',\
           'node_size': 0.5,\
           'edge_color':'gray',\
           'width': 0.2}
with open('../data/subgraphs/twitch_es_graph.json','r') as f:
    data = json.loads(f.read())
g = readwrite.json_graph.node_link_graph(data)

new_users = list(pd.read_csv('../data/subgraphs/twitch_es_new_users.csv').id)
all_users = list(g.nodes)
old_users = [u for u in all_users if u not in new_users]
plt.figure(figsize=(8,5))
pos = nx.spring_layout(g)
nx.draw(g, pos, **options)
# Draw ego as large and red
nx.draw_networkx_nodes(g, pos, nodelist=new_users, node_size=3, node_color='r')
plt.show()

old_graph = g.subgraph(nodes=old_users)
nx.draw(old_graph,**options)

# 顺便生成一下训练node2vec需要的edgelist吧：
with open('../data/subgraphs/%s.edgelist'%name,'w') as f:
    for pair in list(old_graph.edges):
        f.write("%s %s\n"%(pair[0],pair[1]))
#==================
"""
graphsage要求的第一个文件就是-G.json. A networkx-specified json file describing the input graph. 
Nodes have 'val' and 'test' attributes specifying if they are a part of the validation and test sets, respectively.
每个节点的属性包括："test","id","feature"(list),"val","label"(list)

<train_prefix>-id_map.json -- A json-stored dictionary mapping the graph node ids to consecutive integers.
<train_prefix>-class_map.json -- A json-stored dictionary mapping the graph node ids to classes.
"""
# ======== creating G.json ==========

# add nodes and attributes:

#for node in old_users:
#    old_graph.add_nodes_from([(int(node),{"test":False,"val":False,"feature":[],"label":[]})])

# 这里我的graph本来就有了，现在根据graphsage输入格式的的要求，给node添加一些属性
for user in old_users:
    old_graph.node[user]["test"] = False
    old_graph.node[user]["val"] = False
    old_graph.node[user]["feature"] = []
    old_graph.node[user]["label"] = []

## add edges:
#for each in relation_df.iterrows():
#    g.add_edge(int(each[1][0]),int(each[1][1]))

g_dict = json_graph.node_link_data(old_graph)
with open('../data/subgraphs/%s-G.json'%name,'w') as f:
    f.write(json.dumps(g_dict))

#nx.draw(g,with_labels=True)

# ===== creating id_map.json and class_map.json===========
id_map = {}
class_map = {}
for i, id in enumerate(old_users):
    id_map[str(id)] = i
    class_map[str(id)] = []
with open('../data/subgraphs/%s-id_map.json'%name,'w') as f:
    f.write(json.dumps(id_map))
with open('../data/subgraphs/%s-class_map.json'%name,'w') as f:
    f.write(json.dumps(class_map))

#========= creating -walks.txt ===========
import random
def run_random_walks(G, nodes, num_walks):
    pairs = []
    for count, node in enumerate(nodes):
        if G.degree(node) == 0:
            continue
        for i in range(num_walks):
            curr_node = node
            for j in range(3): # walk length
                next_node = random.choice(list(G.neighbors(curr_node)))
                # self co-occurrences are useless
                if curr_node != node:
                    pairs.append((node,curr_node))
                curr_node = next_node
        if count % 1000 == 0:
            print("Done walks for", count, "nodes")
    return pairs

walk_pairs = run_random_walks(old_graph,old_users,50)
with open('../data/subgraphs/%s-walks.txt'%name, "w") as fp:
    fp.write("\n".join([str(p[0]) + "\t" + str(p[1]) for p in walk_pairs]))

# using graphsage to train the embeddings......

#============ transfer the graphsage embdding to word2vec format: =============
""" gensim word2vec format
The output file has n+1 lines for a graph with n vertices. The first line has the following format:
           num_of_nodes dim_of_representation

The next n lines are as follows:
           node_id dim1 dim2 ... dimd
"""
emb = np.load('../data/subgraphs/twitch_es_gcn.npy')

with open('../node_embeddings/twitch_es_gcn.emb','w') as emb_f:
    emb_f.write(str(len(emb))+' 256\n')
    with open('../data/subgraphs/twitch_es_gcn.txt') as order_f:
        for i,line in enumerate(order_f):
            node_name = line.strip()
            node_feat = ' '.join([str(x) for x in emb[i]])
            emb_f.write(node_name+' '+node_feat+'\n')

# using gensim to load the embeddings:
from gensim.models import KeyedVectors
gcn = KeyedVectors.load_word2vec_format("data/subgraphs/subgraph1.emb")



# 实例可视化，看看gcn embedding的直观效果：
center = 8236
friends = list(old_graph.neighbors(center))
arounds = []
for n in list(old_graph.neighbors(center)):
    arounds.append(n)
    arounds += old_graph.neighbors(n)
arounds = list(set(arounds))
ego_graph = old_graph.subgraph(arounds)

similars = [int(x[0]) for x in list(gcn.most_similar(str(center),topn=len(friends)))]
similars = [n for n in similars if n in arounds]

cross_users = list(set(friends).intersection(set(similars)))

plt.figure(figsize=(8,5))
pos = nx.spring_layout(ego_graph)
# 画全图：
nx.draw(ego_graph, pos, **options)
# 画中心点：
nx.draw_networkx_nodes(ego_graph, pos, nodelist=[center], node_size=100, node_color='black', label='Center User')
# 画朋友节点：
nx.draw_networkx_nodes(ego_graph, pos, nodelist=friends, node_size=30, node_color='r', label='Friends')
# 画gcn拓展节点：
nx.draw_networkx_nodes(ego_graph, pos, nodelist=similars, node_size=30, node_color='b', label='GCN Similars')
# 画交叉节点：
nx.draw_networkx_nodes(ego_graph, pos, nodelist=cross_users, node_size=30, node_color='purple', label='Cross Users')
fontdict={'family': 'Arial', 'color': 'black', 'weight': 'normal', 'size': 20}
plt.title("Ego-graph for User %s"%center,fontdict=fontdict)
plt.legend()
plt.show()
















